#include <asm/io.h>
#include <linux/module.h>
#include <linux/pci.h>
#include <linux/delay.h>
#include <linux/pid.h>
#include <linux/kernel.h>
#include <linux/kthread.h>
#include <linux/types.h>
#include <linux/slab.h>
#include <linux/atomic.h>

// 保存edu设备信息
struct edu_dev_info
{
	resource_size_t io;
	long range, flags;
	void __iomem *ioaddr;
	int irq;
};

static struct pci_device_id id_table[] = {
	{PCI_DEVICE(0x1234, 0x11e8)}, // edu设备id
	{0} // 最后一组是0，表示结束
};

struct edu_dev_info *edu_info;
spinlock_t lock;

/// @brief edu设备发现函数
/// @param dev 
/// @param id 
/// @return 
static int edu_driver_probe(struct pci_dev *dev, const struct pci_device_id *id)
{
	int ret = 0;
	printk("executing edu driver probe function!\n");

	ret = pci_enable_device(dev);
	if (ret)
	{
		printk(KERN_ERR "IO Error.\n");
		return -EIO;
	}

	// 为edu_dev_info实例分配内存
	edu_info = kmalloc(sizeof(struct edu_dev_info), GFP_KERNEL);
	if (!edu_info)
	{
		ret = -ENOMEM;
		goto out_mypci;
	}

	// 获取BAR0的总线地址、范围和标志
	edu_info->io = pci_resource_start(dev, 0);
	edu_info->range = pci_resource_len(dev, 0);
	edu_info->flags = pci_resource_flags(dev, 0);
	edu_info->irq = dev->irq; // 从PCI设备结构体获取中断号

	// 将BAR的总线地址映射到系统内存的虚拟地址
	edu_info->ioaddr = pci_ioremap_bar(dev, 0);
	if (!edu_info->ioaddr)
	{
		printk(KERN_ERR "PCI ioremap bar failed!\n");
		ret = -ENOMEM;
		goto out_kfree;
	}

	ret = pci_request_regions(dev, "edu_dirver"); // 申请一块驱动掌管的内存空间
	if (ret)
	{
		printk("PCI request regions err!\n");
		goto out_iounmap;
	}

	pci_set_drvdata(dev, edu_info); // 设置驱动私有数据
	printk("Probe succeeds.PCIE ioport addr start at %llX, edu_info->ioaddr is 0x%p.\n", 
		   (unsigned long long)edu_info->io, edu_info->ioaddr);
	return 0;

out_iounmap:
	iounmap(edu_info->ioaddr);
out_kfree:
	kfree(edu_info);
out_mypci:
	pci_disable_device(dev);
	return ret;
}

/// @brief edu设备移除函数
/// @param dev 
static void edu_driver_remove(struct pci_dev *dev)
{
	// 从设备结构体dev中提取edu_dev_info实例
	struct edu_dev_info *edu_info = pci_get_drvdata(dev);

	// 取消BAR区域的总线地址映射
	iounmap(edu_info->ioaddr);

	pci_release_regions(dev);

	// 释放edu_dev_info实例
	kfree(edu_info);

	pci_disable_device(dev);
	printk("Device is removed successfully.\n");
}

MODULE_DEVICE_TABLE(pci, id_table); // 暴露驱动能发现的设备ID表单
static struct pci_driver pci_driver = {
	.name = "edu_dirver",
	.id_table = id_table,
	.probe = edu_driver_probe,
	.remove = edu_driver_remove,
};

// =============================================================================== //
#define EDU_DEV_MAJOR 200  /* 主设备号 */
#define EDU_DEV_NAME "edu" /* 设备名 */
int current_id = 0;

struct user_data
{
	int id;
	atomic64_t data; // 原子变量存储阶乘结果
};

struct thread_data
{
	struct user_data* user_data_ptr;
	int input_data; // 用户传入的待计算数值
};

static int kthread_handler(void *data)
{
	struct thread_data* thread_data_ptr = (struct thread_data*)data;
	if (!thread_data_ptr)
		return -EINVAL;

	struct user_data *ud = thread_data_ptr->user_data_ptr;
	int input = thread_data_ptr->input_data;
	uint64_t result = 0;
	unsigned long flags;

	printk("ioctl cmd 0 : factorial calculation for %d\n", input);

	// 加自旋锁，保证设备访问的原子性
	spin_lock_irqsave(&lock, flags);
	{
		// 向edu设备偏移0x08的寄存器写入数值，触发阶乘计算
		writel(input, edu_info->ioaddr + 0x08);
		// 读取计算结果
		result = readl(edu_info->ioaddr + 0x08);
	}
	spin_unlock_irqrestore(&lock, flags);

	// 原子操作将结果存入user_data，确保读写原子性
	atomic64_set(&ud->data, result);

	// 释放线程数据（在ioctl中分配）
	kfree(thread_data_ptr);
	return 0;
}

/// @brief open处理函数
/// @param inode 
/// @param filp 
/// @return 
static int edu_dev_open(struct inode *inode, struct file *filp)
{
	struct user_data* user_data_ptr = kmalloc(sizeof(struct user_data), GFP_KERNEL);
	if (!user_data_ptr)
		return -ENOMEM;

	user_data_ptr->id = current_id++;
	atomic64_set(&user_data_ptr->data, 0); // 初始化原子变量为0

	// 绑定用户进程上下文信息到file结构体
	filp->private_data = user_data_ptr;
	printk("edu device opened, user id: %d\n", user_data_ptr->id);
	return 0;
}

/// @brief close处理函数
/// @param inode 
/// @param filp 
/// @return 
static int edu_dev_release(struct inode *inode, struct file *filp)
{
	// 释放open中分配的user_data内存
	struct user_data *user_data_ptr = filp->private_data;
	if (user_data_ptr)
	{
		printk("edu device closed, user id: %d\n", user_data_ptr->id);
		kfree(user_data_ptr);
		filp->private_data = NULL;
	}
	return 0;
}

/// @brief ioctl处理函数
/// @param pfilp_t 
/// @param cmd 
/// @param arg 
/// @return 
static long edu_dev_unlocked_ioctl(struct file *pfilp_t, unsigned int cmd, unsigned long arg)
{
	struct user_data *user_data_ptr = pfilp_t->private_data;
	if (!user_data_ptr)
		return -EINVAL;

	switch (cmd)
	{
	case 0: // cmd=0：传入待计算的数值，创建内核线程计算
	{
		// 分配线程数据结构
		struct thread_data *td = kmalloc(sizeof(struct thread_data), GFP_KERNEL);
		if (!td)
			return -ENOMEM;

		td->user_data_ptr = user_data_ptr;
		td->input_data = (int)arg; // 转换用户传入的参数

		// 创建内核线程
		struct task_struct *task = kthread_create(kthread_handler, td, "edu_factorial_thread");
		if (IS_ERR(task))
		{
			kfree(td);
			return PTR_ERR(task);
		}

		// 唤醒内核线程开始执行
		wake_up_process(task);
		printk("created factorial thread for input: %d\n", td->input_data);
		return 0;
	}
	case 1: // cmd=1：读取阶乘计算结果
	{
		uint64_t res = atomic64_read(&user_data_ptr->data);
		printk("read factorial result: %llu\n", (unsigned long long)res);
		return (long)res;
	}
	default:
		printk("invalid ioctl cmd: %u\n", cmd);
		return -EINVAL;
	}
}

static struct file_operations edu_dev_fops = {
	.open = edu_dev_open,
	.release = edu_dev_release,
	.unlocked_ioctl = edu_dev_unlocked_ioctl,
};

/// @brief 驱动程序初始化
/// @param  
/// @return 
static int __init edu_dirver_init(void)
{
	printk("HELLO PCI\n");
	int ret = 0;

	// 注册字符设备
	ret = register_chrdev(EDU_DEV_MAJOR, EDU_DEV_NAME, &edu_dev_fops);
	if (0 > ret)
	{
		printk("kernel edu dev register_chrdev failure\n");
		return -1;
	}
	printk("chrdev edu dev is insmod, major_dev is 200\n");

	// 注册edu pci设备
	ret = pci_register_driver(&pci_driver);
	if (ret)
	{
		printk("kernel edu dev pci_register_driver failure\n");
		unregister_chrdev(EDU_DEV_MAJOR, EDU_DEV_NAME);
		return ret;
	}

	// 初始化自旋锁
	spin_lock_init(&lock);
	return 0;
}

/// @brief 驱动程序注销
/// @param  
/// @return 
static void __exit edu_dirver_exit(void)
{
	// 注销字符设备
	unregister_chrdev(EDU_DEV_MAJOR, EDU_DEV_NAME);
	// 注销edu pci设备
	pci_unregister_driver(&pci_driver);
	printk("GOODBYE PCI\n");
}

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Luo");
MODULE_DESCRIPTION("homework3");
MODULE_VERSION("1.0");

module_init(edu_dirver_init);
module_exit(edu_dirver_exit);